from vtkplotter import *
import vtk
import argparse
import numpy as np
import os
import sys
import json
import math
import time
from vtk.util import numpy_support
from scipy.optimize import minimize
from scipy.stats import f

scriptdir = os.path.dirname(sys.argv[0])
default_head = os.path.join(scriptdir, "head/head.vtk")
default_diean = os.path.join(scriptdir, "head/diean.vtk")
parser = argparse.ArgumentParser()
parser.add_argument('--head', type=str, default=default_head,
                    help='path of the head model')
parser.add_argument('--diean', type=str, default=default_diean,
                    help='path of the diean model')
parser.add_argument('--npsline', type=int, default=2,
                    help='number of control points for a spline')
parser.add_argument('--res', type=float, default=10,
                    help='resolution of the spline')

args = parser.parse_args()


def normalize(vec):
    d = np.linalg.norm(vec)
    if d > 0:
        return vec / d
    return vec


def optimize_vmin_vmax(values):
    vmin = values.min()
    vmax = values.max()
    mean = values.mean()
    std = values.std()
    std = math.sqrt(std)
    vmin = max(vmin, mean - std)
    vmax = min(vmax, mean + std)
    return vmin, vmax


def frenet_serret(pts):
    tangents = np.zeros(pts.shape)
    normals = np.zeros(pts.shape)
    binormals = np.zeros(pts.shape)
    frenets = np.zeros(pts.shape)
    for i in range(pts.shape[0] - 1):
        tangents[i] = pts[i + 1] - pts[i]
        ds = np.linalg.norm(tangents[i])
        frenets[i + 1, 0] = frenets[i, 0] + ds  # length
        if ds > 0:
            tangents[i] = tangents[i] / ds
            if i > 0:
                tangent_2nd = normalize(tangents[i] - tangents[i - 1])
                binormals[i] = normalize(np.cross(tangents[i], tangent_2nd))
                normals[i] = normalize(np.cross(binormals[i], tangents[i]))
                frenets[i, 1] = math.acos(
                    np.clip(np.abs(np.dot(tangents[i - 1], tangents[i])), 0, 1)) / ds  # curvature
                if i > 1:
                    frenets[i, 2] = math.acos(
                        np.clip(np.abs(np.dot(binormals[i - 1], binormals[i])), 0, 1)) / ds  # torsion

    frenets[-1, 1] = frenets[-2, 1]
    frenets[-1, 2] = frenets[-2, 2]
    return frenets, tangents, normals, binormals


def compute_pca_coordinates(points, camera, weights=None):
    cov = np.cov(points, rowvar=0, aweights=weights)     # covariance matrix
    U, s, R = np.linalg.svd(cov)  # singular value decomposition
    p, n = s.size, points.shape[0]
    pvalue = 0.95
    fppf = f.ppf(pvalue, p, n-p)*(n-1)*p*(n+1)/n/(n-p)  # f % point function
    ua, ub, uc = np.sqrt(s*fppf)*2  # semi-axes (largest first)
    center = np.mean(points, axis=0)     # centroid of the hyperellipsoid
    sphericity = (((ua-ub)/(ua+ub))**2
                  + ((ua-uc)/(ua+uc))**2
                  + ((ub-uc)/(ub+uc))**2)/3. * 4.

    matri = vtk.vtkMatrix4x4()
    matri.DeepCopy((R[0][0] * ua, R[1][0] * ub, R[2][0] * uc, center[0],
                    R[0][1] * ua, R[1][1] * ub, R[2][1] * uc, center[1],
                    R[0][2] * ua, R[1][2] * ub, R[2][2] * uc, center[2], 0, 0, 0, 1))
    vtra = vtk.vtkTransform()
    vtra.SetMatrix(matri)

    axes = []
    for ax in ([1, 0, 0], [0, 1, 0], [0, 0, 1]):
        l = vtk.vtkLineSource()
        l.SetPoint1([0, 0, 0])
        l.SetPoint2(ax)
        l.Update()
        t = vtk.vtkTransformFilter()
        t.SetTransform(vtra)
        t.SetInputData(l.GetOutput())
        t.Update()
        axes.append(t.GetOutput())

    points = []
    for axis in axes:
        pts = numpy_support.vtk_to_numpy(axis.GetPoints().GetData())
        origin = np.array(pts[0])
        points.append(pts[1])

    camera_pos = np.array(camera.GetPosition())
    camera_target = np.array(camera.GetFocalPoint())
    camera_up = np.array(camera.GetViewUp())
    camera_forward = normalize(camera_pos - camera_target)
    front = origin
    up = origin
    max_front_dot = 0
    max_up_dot = 0
    for pt in points:
        vec = pt - origin
        l = np.linalg.norm(vec)
        vec = vec / l
        dot = abs(np.dot(camera_forward, vec))
        if dot > max_front_dot:
            max_front_dot = dot
            front = vec
        dot = abs(np.dot(camera_up, vec))
        if dot > max_up_dot:
            max_up_dot = dot
            up = vec

    left = np.cross(up, front)
    return origin, left, up, front


def model_symmetry(x, origin, up, points):
    left = normalize(np.array((x[1], x[2], x[3])))
    front = normalize(np.cross(up, left))
    # origin = args[0]
    # points = args[1]
    origin = origin + left * x[0]
    points = points - origin
    points = points[np.dot(points, front) > 0]
    dot = np.dot(points, left)
    left_side = points[dot > 0]
    left_center = left_side.mean(axis=0)
    right_side = points[dot < 0]
    right_center = right_side.mean(axis=0)
    return np.dot(left_center + right_center - 2 * origin, left)


class SymmetryApp(object):
    def __init__(self, args):
        self.args = args
        self.cmap = 'jet'
        self.vp = Plotter()

        self.head_model = self.vp.load(args.head, c='gray', alpha=0.5)
        self.simplied_head_model = None
        self.diean_model = self.vp.load(args.diean, c='orangered', alpha=0.5)
        self.working_dir = os.path.dirname(args.head)
        self.current_mode = None
        self.build_stage = None
        self.current_actor = None
        self.current_point_index = 0

        self.NPointOfSpline = args.npsline
        self.left_spline_pts = Points(
            np.zeros((self.NPointOfSpline, 3)), r=12, c='magenta')
        self.right_spline_pts = Points(
            np.zeros((self.NPointOfSpline, 3)), r=12, c='navy')
        # coordinate_points = Points(np.zeros((3, 5)), c=['red', 'green', 'blue', 'cyan', 'yellow', 'orange'])
        self.N = Point(np.zeros(3), c='red')
        self.S = Point(np.zeros(3), c='green')
        self.Or_left = Point(np.zeros(3), c='blue')
        self.Or_right = Point(np.zeros(3), c='cyan')
        self.Po_left = Point(np.zeros(3), c='yellow')
        self.Po_right = Point(np.zeros(3), c='orange')
        self.S.SetPosition(self.diean_model.centerOfMass())
        self.vp.add(self.S)

        self.left_spline = None
        self.right_spline = None
        self.mirror_spline = None
        self.left_distances = None
        self.right_distances = None
        self.left_frenets = None
        self.right_frenets = None
        self.mirror_frenets = None
        self.color_mode = 'symmetry'
        self.silhouette = None
        self.mirror_plane = None

    def save(self):
        with open(os.path.join(self.working_dir, 'result.txt'), 'w') as f:
            data = object()
            data.N = self.N.getPoints()[0]
            data.S = self.S.getPoints()[0]
            data.Or_left = self.Or_left.getPoints()[0]
            data.Or_right = self.Or_right.getPoints()[0]
            data.Po_left = self.Po_left.getPoints()[0]
            data.Po_right = self.Po_right.getPoints()[0]
            data.left_spline_pts = self.left_spline_pts.GetPoints()
            data.right_spline_pts = self.right_spline_pts.GetPoints()
            if self.build_stage == 'complete':
                data.left_spline = self.left_spline.GetPoints()
                data.right_spline = self.right_spline.GetPoints()
                data.left_distances = self.left_distances
                data.right_distances = self.right_distances
            json.dump(data, f)

    def build_coordinates(self):
        if self.Or_left not in self.vp.getActors():
            printc("no feature point, skip building coordinates.")
            return

        printc("build coordinates.")
        hpts = np.zeros((4, 3))
        hpts[0] = self.Or_left.getPoints()
        hpts[1] = self.Or_right.getPoints()
        hpts[2] = self.Po_left.getPoints()
        hpts[3] = self.Po_right.getPoints()
        horizontal_plane = fitPlane(hpts)
        up = horizontal_plane.info['normal']
        origin = self.N.getPoints()[0]
        if np.dot(origin, up) < 0:
            up = -up

        front = origin - self.S.getPoints()[0]
        front = front - up * np.dot(front, up)
        left = np.cross(up, front)
        self.set_coordinates(origin, left, front, up)

    def set_coordinates(self, origin, left, front, up):
        up = normalize(up)
        front = normalize(front)
        left = normalize(left)
        maxSize = self.head_model.maxBoundSize()
        if self.mirror_plane and self.mirror_plane in self.vp.getActors():
            self.vp.remove(self.mirror_plane)
        self.mirror_plane = Plane(
            pos=origin, normal=left, sx=maxSize, sy=maxSize, c='yellow', alpha=0.3)
        self.vp.add(self.mirror_plane)
        self.coord_origin = origin
        self.coord_left = left
        self.coord_front = front
        self.coord_up = up

    def build_spline(self, points, silhouette, c):
        ptIdx = []
        ptIdx.append(silhouette.closestPoint(
            points.getPoint(0), N=2, returnIds=True)[0])
        closest = silhouette.getPoint(ptIdx[0])
        points.setPoint(0, closest)
        ptIdx.append(silhouette.closestPoint(
            points.getPoint(1), N=2, returnIds=True)[0])
        closest = silhouette.getPoint(ptIdx[1])
        points.setPoint(1, closest)

        try:
            path = self.find_shortest_path(silhouette, ptIdx[0], ptIdx[1])
            return spline(silhouette.getPoints()[path], smooth=0.5, s=5, res=1).color(c)
        except Exception:
            return silhouette.color(c)

    def find_shortest_path(self, silhouette, start, end):
        poly = silhouette.polydata()
        import networkx as nx
        g = nx.Graph()
        lines = poly.GetLines().GetData()
        for i in range(poly.GetNumberOfCells()):
            pointIdList = vtk.vtkIdList()
            poly.GetCellPoints(i, pointIdList)
            g.add_edge(pointIdList.GetId(0), pointIdList.GetId(1))

        path = nx.shortest_path(g, source=start, target=end)
        return path

    def detect_mirror_plane(self, tol=1e-4):
        # front_silhouette = self.head_model.silhouette(self.vp.camera)
        # self.vp.add(front_silhouette)
        # points = front_silhouette.getPoints()
        points = self.head_model.getPoints()
        origin, left, up, front = compute_pca_coordinates(
            points, self.vp.camera)
        bounds = [(-self.head_model.diagonalSize() / 2,
                   self.head_model.diagonalSize() / 2), (-1, 1), (-1, 1), (-1, 1)]
        res = minimize(model_symmetry, (0, left[0], left[1], left[2]), args=(
            origin, up, points), bounds=bounds)
        left = normalize(np.array((res.x[1], res.x[2], res.x[3])))
        front = normalize(np.array((res.x[3], res.x[4], res.x[5])))
        up = normalize(np.cross(front, left))
        origin = origin + left * res.x[0]
        self.set_coordinates(origin, left, front, up)
        self.vp.interactor.Render()
        # iter = 0
        # while iter < 10:
        #     old_origin, old_left, old_up, old_front = origin, left, up, front
        #     weights = np.dot(points - origin, front)
        #     weights[weights < 0] = 0
        #     # weights = weights * weights
        #     weights = (weights - weights.min()) / (weights.max() - weights.min())
        #     origin, left, up, front = compute_pca_coordinates(points, self.vp.camera, weights=weights)
        #     tvalue = abs(np.sum(np.dot(points - origin, left)))
        #     iter = iter+1
        #     self.set_coordinates(origin, left, front, up)
        #     self.vp.interactor.Render()
        #     time.sleep(0.5)

    def mirror(self):
        printc("do the mirror ...")
        mirror_pts = self.left_spline.getPoints(copy=True)
        right_pts = self.right_spline.getPoints()

        # do the mirror
        for i in range(mirror_pts.shape[0]):
            d = mirror_pts[i] - self.coord_origin
            mirror_pts[i] = -self.coord_left * np.dot(d, self.coord_left) + self.coord_up * np.dot(
                d, self.coord_up) + self.coord_front * np.dot(d, self.coord_front) + self.coord_origin

        self.mirror_spline = spline(
            mirror_pts, smooth=0, s=5, res=1).color('magenta')
        self.vp.add(self.mirror_spline)
        self.left_frenets, _, _, _ = frenet_serret(
            self.left_spline.getPoints())
        self.right_frenets, _, _, _ = frenet_serret(
            self.right_spline.getPoints())
        self.mirror_frenets, _, _, _ = frenet_serret(
            self.mirror_spline.getPoints())
        self.toggle_color_scalars('torsion')
        printc("mirror completed.")

    def compute_symmetry(self):
        printc("evaluate symmetry ...")
        mirror_pts = self.mirror_spline.getPoints()
        printc('number of mirror points: %s' % mirror_pts.shape[0])
        right_pts = self.right_spline.getPoints()
        printc('number of right points: %s' % right_pts.shape[0])
        total = mirror_pts.shape[0] + right_pts.shape[0]
        step = total // 10
        printc('total points: %s, progress report step: %s' % (total, step))
        self.left_distances = np.zeros(mirror_pts.shape[0])
        self.right_distances = np.zeros(right_pts.shape[0])
        pb = ProgressBar(0, 100, char='evaluate symmetry')
        for i in range(mirror_pts.shape[0]):
            closest = self.right_spline.closestPoint(mirror_pts[i])
            self.left_distances[i] = np.linalg.norm(mirror_pts[i] - closest)
            if i % step == 0:
                pb.print(counts=(100 * i // total))

        for i in range(right_pts.shape[0]):
            closest = self.mirror_spline.closestPoint(right_pts[i])
            self.right_distances[i] = np.linalg.norm(right_pts[i] - closest)
            if i % step == 0:
                pb.print(counts=(100 * (i + mirror_pts.shape[0]) // total))

        pb.print(counts=100)
        printc("colorize symmetry of splines ...")
        self.toggle_color_scalars('symmetry')
        left_length = np.sum(self.left_frenets[:, 0])
        right_length = np.sum(self.right_frenets[:, 0])
        symmetry_index = (left_length + right_length) / (self.left_distances.std() +
                                                         self.right_distances.std() + (left_length + right_length) / 1000) / 1000
        printc('symmetry index is %s, evaluation completed.' % symmetry_index)
        self.vp.add(Text('Symmetry Index: %s' % symmetry_index))

    def toggle_color_scalars(self, mode=None):
        if mode == None:
            if self.color_mode == 'symmetry':
                mode = 'torsion'
            elif self.color_mode == 'torsion':
                mode = 'curvature'
            elif self.color_mode == 'curvature':
                mode = 'symmetry'

        self.color_mode = mode
        if self.mirror_spline == None:
            return

        if self.color_mode == 'symmetry':
            if self.left_distances is None:
                return

            vmin = min(np.min(self.left_distances),
                       np.min(self.right_distances))
            vmax = max(np.max(self.left_distances),
                       np.max(self.right_distances))
            self.left_spline.pointColors(
                self.left_distances, cmap=self.cmap, vmin=vmin, vmax=vmax)
            self.mirror_spline.pointColors(
                self.left_distances, cmap=self.cmap, vmin=vmin, vmax=vmax)
            self.right_spline.pointColors(
                self.right_distances, cmap=self.cmap, vmin=vmin, vmax=vmax)
        else:
            idx = -1
            if self.color_mode == 'torsion':
                idx = 2
            elif self.color_mode == 'curvature':
                idx = 1

            if idx > 0:
                vmin, vmax = optimize_vmin_vmax(self.left_frenets[:, idx])
                self.left_spline.pointColors(
                    self.left_frenets[:, idx], cmap=self.cmap, vmin=vmin, vmax=vmax)
                vmin, vmax = optimize_vmin_vmax(self.mirror_frenets[:, idx])
                self.mirror_spline.pointColors(
                    self.mirror_frenets[:, idx], cmap=self.cmap, vmin=vmin, vmax=vmax)
                vmin, vmax = optimize_vmin_vmax(self.right_frenets[:, idx])
                self.right_spline.pointColors(
                    self.right_frenets[:, idx], cmap=self.cmap, vmin=vmin, vmax=vmax)

        self.vp.interactor.Render()


def onKeyPress(key):
    global app
    key = key.lower()
    if key == "a":
        printc("toggle transparency ...")
        if app.head_model.alpha() == 1.0:
            app.head_model.alpha(0.5)
        else:
            app.head_model.alpha(1.0)

        app.vp.interactor.Render()
        return

    if key == "n":
        printc("pick a point for N ...")
        app.current_mode = 'N'
        app.current_actor = app.N
        return

    if key == "s":
        printc("save to result.txt")
        app.save()
        return

    if key == "o":
        printc("press L or R to ensure left or right ide of Or to pick ...")
        app.current_mode = 'Or'
        app.current_actor = None
        return

    if key == "l" and app.current_mode == 'Or':
        printc("pick a point for Or_left ...")
        app.current_mode = 'Or_left'
        app.current_actor = app.Or_left
        return

    if key == "r" and app.current_mode == 'Or':
        printc("pick a point for Or_right ...")
        app.current_mode = 'Or_right'
        app.current_actor = app.Or_right
        return

    if key == "d":
        printc("press L or R to ensure left or right side of Po to pick ...")
        app.current_mode = 'Po'
        app.current_actor = None
        return

    if key == "l" and app.current_mode == 'Po':
        printc("pick a point for Po_left ...")
        app.current_mode = 'Po_left'
        app.current_actor = app.Po_left
        return

    if key == "r" and app.current_mode == 'Po':
        printc("pick a point for Po_right ...")
        app.current_mode = 'Po_right'
        app.current_actor = app.Po_right
        return

    if key == "e":
        printc("press L or R to ensure left or right side of spline to pick ...")
        app.current_mode = 'Sp'
        app.current_actor = None
        app.current_point_index = 0
        return

    if key == "l" and app.current_mode == 'Sp':
        printc("pick a point for Spline_left ...")
        app.current_mode = 'Sp_left'
        app.current_actor = app.left_spline_pts
        return

    if key == "r" and app.current_mode == 'Sp':
        printc("pick a point for Spline_right ...")
        app.current_mode = 'Sp_right'
        app.current_actor = app.right_spline_pts
        return

    if key.isdigit() and app.current_mode.startswith('Sp_'):
        app.current_point_index = int(key)
        app.current_point_index = min(
            app.current_point_index, app.NPointOfSpline - 1)
        printc("pick %sth point for %s ..." %
               (app.current_point_index, app.current_mode))
        return

    if key == 'b' and app.build_stage == None:
        app.build_stage = 'coordinates'
        app.build_coordinates()
        return

    if key == 'b' and app.build_stage == 'coordinates':
        app.build_stage = 'mirror'
        app.mirror()
        return

    if key == 'b' and app.build_stage == 'mirror':
        app.build_stage = 'symmetry'
        app.compute_symmetry()
        app.build_stage = 'complete'
        return

    if key == 'c':
        app.toggle_color_scalars()
        return

    if key == 'x':
        app.detect_mirror_plane()
        return

    app.current_mode = None


def clean_silhouette(silhouette):
    poly = silhouette.polydata()
    linked_points_dict = {}
    linked_points_ids = []
    new_cells = numpy_support.vtk_to_numpy(poly.GetLines().GetData()).copy()
    cell_pos = 0
    for i in range(poly.GetNumberOfCells()):
        cell_size = new_cells[cell_pos]
        cell_pos = cell_pos + 1
        for j in range(cell_size):
            pid = new_cells[cell_pos]
            if pid not in linked_points_dict:
                linked_points_dict[pid] = len(linked_points_ids)
                linked_points_ids.append(pid)
            new_cells[cell_pos] = linked_points_dict[pid]
            cell_pos = cell_pos + 1

    new_pts = silhouette.getPoints()[linked_points_ids].copy()

    ppoints = vtk.vtkPoints()  # Generate the polyline for the spline
    profileData = vtk.vtkPolyData()
    ppoints.SetData(numpy_support.numpy_to_vtk(new_pts))
    poly.SetPoints(ppoints)

    lines = vtk.vtkCellArray()
    lines.SetCells(poly.GetNumberOfCells(), numpy_support.numpy_to_vtkIdTypeArray(
        new_cells))  # Create the polyline
    poly.SetLines(lines)
    silhouette.updateMesh(poly)
    return silhouette


def onRightClick(actor):
    global app
    if app.current_mode is not None and app.current_actor is not None:
        if app.current_mode.startswith('Sp_'):
            app.current_point_index = min(
                app.current_point_index, app.NPointOfSpline - 1)
            printc("pick %sth point for %s ..." %
                   (app.current_point_index, app.current_mode))
            app.current_actor.setPoint(
                app.current_point_index, app.vp.picked3d)
            if app.current_point_index == 0:
                if app.current_mode == 'Sp_left':
                    app.left_silhouette = clean_silhouette(
                        app.head_model.silhouette(app.vp.camera))
                elif app.current_mode == 'Sp_right':
                    app.right_silhouette = clean_silhouette(
                        app.head_model.silhouette(app.vp.camera))
            else:
                printc("build splines ...")
                if app.current_mode == 'Sp_left':
                    if app.left_spline and app.left_spline in app.vp.getActors():
                        app.vp.remove(app.left_spline)
                    app.left_spline = app.build_spline(
                        app.left_spline_pts, app.left_silhouette, 'magenta')
                    app.vp.add(app.left_spline)
                elif app.current_mode == 'Sp_right':
                    if app.right_spline and app.right_spline in app.vp.getActors():
                        app.vp.remove(app.right_spline)
                    app.right_spline = app.build_spline(
                        app.right_spline_pts, app.right_silhouette, 'navy')
                    app.vp.add(app.right_spline)

                app.vp.interactor.Render()

            app.current_point_index = app.current_point_index + 1
        else:
            printc("position", [app.vp.picked3d], c="g")
            app.current_actor.SetPosition(app.vp.picked3d)

        actor_added = False
        for actor in app.vp.getActors():
            if actor == app.current_actor:
                actor_added = True
                break

        if not actor_added:
            app.vp.add(app.current_actor)


app = SymmetryApp(args)
app.vp.keyPressFunction = onKeyPress
app.vp.mouseRightClickFunction = onRightClick
app.vp.show(interactive=True)
